#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from nnunet.training.network_training.competitions_with_custom_Trainers.BraTS2020.nnUNetTrainerV2BraTSRegions_moreDA import (
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_largeUnet_Groupnorm,
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_axialattention_unet,
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet,
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_normalUnet_axialattention,
)
from nnunet.network_architecture.initialization import InitWeights_He
from nnunet.network_architecture.yunet import YUNet
# from nnunet.network_architecture.hecktor import HecktorNet
import torch
from torch import nn
import contextlib


class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_largeUnet_Groupnorm_yusongli(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_largeUnet_Groupnorm
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 10

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 14  # default: 30


# class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yusongli(nnUNetTrainerV2BraTSRegions_DA4_BN_BD):
#     def __init__(self, *args: object, **kwargs: object) -> None:
#         super().__init__(*args, **kwargs)
#         self.max_num_epochs = 10


class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_axialattention_unet_yusongli(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_axialattention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000
        self.volume_shape = (24, 64, 80)

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 15  # default: 30


class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_yusongli(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000
        self.volume_shape = (24, 64, 80)

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 22  # default: 30


class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_normalUnet_axialattention_yusongli(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_normalUnet_axialattention
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000
        self.volume_shape = (24, 64, 80)

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 11  # default: 30


class nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000
        self.volume_shape = (24, 64, 80)

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 12  # default: 30

    def initialize_network(self):
        """inference_apply_nonlin to sigmoid + larger unet"""
        if self.threeD:
            conv_op = nn.Conv3d
            dropout_op = nn.Dropout3d
            norm_op = nn.BatchNorm3d

        else:
            conv_op = nn.Conv2d
            dropout_op = nn.Dropout2d
            norm_op = nn.BatchNorm2d

        norm_op_kwargs = {'eps': 1e-5, 'affine': True}
        dropout_op_kwargs = {'p': 0, 'inplace': True}
        net_nonlin = nn.LeakyReLU
        net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True}
        self.network = YUNet(
            self.num_input_channels,
            self.base_num_features,
            self.num_classes,
            len(self.net_num_pool_op_kernel_sizes),
            self.conv_per_stage,
            2,
            conv_op,
            norm_op,
            norm_op_kwargs,
            dropout_op,
            dropout_op_kwargs,
            net_nonlin,
            net_nonlin_kwargs,
            True,
            False,
            lambda x: x,
            InitWeights_He(1e-2),
            self.net_num_pool_op_kernel_sizes,
            self.net_conv_kernel_sizes,
            False,
            True,
            True,
            320,
            encoder_scale=1,
            axial_attention=True,
            heads=1,
            dim_heads=4,
            volume_shape=self.volume_shape,
            no_attention=[0],
            axial_bn=True,
            sum_axial_out=True,
            residual_attention=True,
        )
        print(self.network)
        if torch.cuda.is_available():
            self.network.cuda()
        self.network.inference_apply_nonlin = nn.Sigmoid()


class hecktor(nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet):
    def initialize_network(self):
        self.network = HecktorNet(
            in_channels=1,
            n_cls=2,
            n_filters=4,
            reduction=2,
        )
        if torch.cuda.is_available():
            self.network.cuda()


# =================
# === Ablations ===
# =================

# 1. Arch: ATT
class yunet_arch_ATT(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 12  # default: 30


# 2. Arch: RCAB
class yunet_arch_RCAB(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 12  # default: 30


# 3. Arch: RCAB + ATT
class yunet_arch_RCABATT(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 10  # default: 30


# 4. Arch: ATT + RCAB
class yunet_arch_ATTRCAB(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 10  # default: 30


# 5. Arch: residual RCAB + ATT no1st
class yunet_arch_resRCABATT(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 10  # default: 30


# 6. Bottleneck + Arch: residual ( RCAB + ATT ) have1st
class yunet_neckarch_resRCABATT_havefirst(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 7  # default: 30


# 7. Bottleneck + Arch: residual ( residual RCAB + residual ATT ) no first layer
class yunet_neckarch_ArchBlock_nofirst(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 10  # default: 30


# 8. Bottleneck + Arch: residual ( residual RCAB + residual ATT ) have first layer
class yunet_neckarch_ArchBlock_havefirst(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 7  # default: 30


# 9. Rollback to most higher dice model
class yunet_arch_resRCABATT_rollback(yunet_arch_resRCABATT):
    pass


# 10. Bottleneck + Arch: residual ( RCAB + ATT ) have1st
class yunet_arch_resRCABATT_havefirst(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 7  # default: 30


# 11. Change
class yunet_change(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class yunet_change2(
    yunet_change
):
    '''
    Concat two tensors, instead of directly add them.
    '''
    pass


class yunet_change3(
    yunet_change
):
    '''
    Concat two tensors and Add residual connections.
    '''
    pass


class yunet_arch_resRCABATT_change2_havefirst(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class Brats21Lab_resaxisatt_batch_06_epoch_300(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300
        self.volume_shape = (24, 64, 80)

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class yunet_arch_resRCABATT_change2_nofirst_bmax_e300(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 20  # default: 30


class yunet_arch_resRCABATT_change2_havefirst_bmax_e300(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 14  # default: 30


class bottleneck_change_00(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_02(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_03(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_04(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_05(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_06(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_change_07(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_00(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_02(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_03(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_04(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_05(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class bottleneck_RCAB_change_06(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


# ===============
# === cluster ===
# ===============
class cluster_yunet_00(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_02(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_03(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_04(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_05(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_06(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_yunet_07(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_brats_00(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_brats_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_brats_02(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_brats_03(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


class cluster_brats_04(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        self.plans['plans_per_stage'][0]['batch_size'] = 6  # default: 30


# ==================
# === caibaoshuo ===
# ==================
class new_fine_yunet(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 300

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_yunet_epoch500(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_yunet_epoch500_01(new_fine_yunet_epoch500):
    pass


class new_fine_yunet_epoch500_02(new_fine_yunet_epoch500):
    pass


class new_fine_yunet_epoch500_03(new_fine_yunet_epoch500):
    pass


class new_fine_yunet_epoch500_04(new_fine_yunet_epoch500):
    pass


class new_fine_yunet_epoch500_05(new_fine_yunet_epoch500):
    pass


class new_fine_yunet_epoch500_06(new_fine_yunet_epoch500):
    pass


class new_fine_brats_epoch500(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 500

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats_epoch500_01(new_fine_brats_epoch500):
    pass


class new_fine_brats_epoch500_02(new_fine_brats_epoch500):
    pass


class new_fine_brats_epoch500_03(new_fine_brats_epoch500):
    pass


class new_fine_brats_epoch500_04(new_fine_brats_epoch500):
    pass


class new_fine_brats_epoch500_05(new_fine_brats_epoch500):
    pass


class new_fine_brats_epoch500_06(new_fine_brats_epoch500):
    pass


class new_fine_yunet_epoch1000(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats_epoch1000(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_yunet_epoch1000_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats_epoch1000_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 1000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_yunet_epoch2000(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_yunet_epoch2000_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_yunet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats_epoch2000(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size


class new_fine_brats_epoch2000_01(
    nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet
):
    def __init__(self, *args: object, **kwargs: object) -> None:
        super().__init__(*args, **kwargs)
        self.max_num_epochs = 2000

    def load_plans_file(self):
        super().load_plans_file()
        batch_size = 4
        self.plans['plans_per_stage'][0]['batch_size'] = batch_size
        with contextlib.suppress(Exception):
            self.plans['plans_per_stage'][1]['batch_size'] = batch_size
